import pandas as pd
import json
import numpy as np
import os


class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)



def dump_dic(folder="multitask"):
    dic = {
        'nodes': [],
        'edges': []
    }
    module = pd.read_csv(os.path.join(folder, "module.csv"))
    module[['hash', 'level']]

    for i in range(len(module)):
        dic['nodes'].append({
            'id': module.loc[i, 'name'],
            'nodeType': "Module",
            'data': module.loc[i, 'property'],
            'hash': module.loc[i, 'hash'],
            'level': module.loc[i, 'level']
        })


    construct = pd.read_csv(os.path.join(folder, "construct.csv"))

    for i in range(len(construct)):
        dic['edges'].append({
            'edgeType': 'construct',
            'source': construct.loc[i,'src'],
            'target': construct.loc[i,'dest'],
            'order': construct.loc[i,'order']
        })


    with open(os.path.join(folder, 'output.json'), 'w') as f:
        json.dump(dic, f, cls=NpEncoder)


def merge(folders):
    out_dic = {
        'nodes': [],
        'edges': []
    }
    ns = set()
    es = set()
    for folder in folders:
        with open(os.path.join(folder, 'output.json'), 'r') as f:
            dic = json.load(f)
        for n in dic['nodes']:
            if n['id'] not in ns:
                out_dic['nodes'].append(n)
                ns.add(n['id'])
        for e in dic['edges']:
            if e['source'] + e['target'] not in es:
                es.add(e['source'] + e['target'])
                out_dic['edges'].append(e)
    return out_dic


if __name__ == "__main__":
    folders = ['mb15', 'mtb_profile', 'multitask', 'kg']
    for folder in folders:
        dump_dic(folder)

    dic = merge(folders)
    with open(os.path.join('output', 'output.json'), 'w') as f:
        json.dump(dic, f, cls=NpEncoder)
